import cv2
import numpy as np
from multiprocessing import Process, Queue
import time


def torch_worker(model_path, input_queue, output_queue):
    """子进程：运行torch模型"""
    import torch  # 在子进程内导入
    from ultralytics import YOLO  # 在子进程内导入

    model = YOLO(model_path)

    while True:
        task = input_queue.get()
        if task is None:  # 终止信号
            break

        image_path = task
        image = cv2.imread(image_path)
        results = model.predict(image, imgsz=640, conf=0.5)
        boxes = [box.xyxy.cpu().numpy() for box in results[0].boxes]
        output_queue.put(boxes)


class PlateDetectionService:
    def __init__(self, yolo_model_path):
        self.input_queue = Queue()
        self.output_queue = Queue()

        # 启动子进程（不传递任何Paddle/不可pickle对象）
        self.process = Process(
            target=torch_worker,
            args=(yolo_model_path, self.input_queue, self.output_queue)
        )
        self.process.start()

        # 主进程初始化PaddleOCR
        from PlateRecognizer import PlateRecognizer
        self.recognizer = PlateRecognizer()

    def process_image(self, image_path):
        """主流程：发送任务到子进程并获取结果"""
        self.input_queue.put(image_path)
        boxes = self.output_queue.get()

        image = cv2.imread(image_path)
        results = []
        for box in boxes[0]:  # 处理第一个检测结果的boxes
            x1, y1, x2, y2 = map(int, box)
            plate_img = image[y1:y2, x1:x2]

            # # 保存plate_img结果到文件
            # cv2.imwrite(f"plate_{time.time()}.jpg", plate_img)

            text, confidence = self.recognizer.recognize_plate(plate_img)
            results.append((text, confidence))

        return results

    def __del__(self):
        """清理子进程"""
        self.input_queue.put(None)
        self.process.join()


if __name__ == "__main__":
    # 示例用法
    detector = PlateDetectionService("../Yolov8_PlateDetect.pt")
    results = detector.process_image("test.jpg")

    for text, conf in results:
        print(f"识别结果: {text} (置信度: {conf:.2f})")
